{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1dfae479-9399-492d-acaa-d9751615ee86",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Finetuning a language model\n",
    "\n",
    "<!--- @wandbcode{dlai_05} -->\n",
    "\n",
    "Let's see how to finetune a language model to generate character backstories using HuggingFace Trainer with wandb integration. We'll use a tiny language model (`TinyStories-33M`) due to resource constraints, but the lessons you learn here should be applicable to large models too!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f0e67f",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoModelForCausalLM\n",
    "from transformers import Trainer, TrainingArguments\n",
    "import transformers\n",
    "transformers.set_seed(42)\n",
    "\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f79c25e3-5f18-4457-84e1-ed2c0d262222",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "wandb.login(anonymous=\"allow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2286ae41-213d-480d-a4ba-8c4e2e1c4771",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"roneneldan/TinyStories-33M\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd80268-c4a1-4e1a-aed3-cd5c3ab4d48f",
   "metadata": {},
   "source": [
    "### Preparing data\n",
    "\n",
    "We'll start by loading a dataset containing Dungeons and Dragons character biographies from Huggingface. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9288a8e-b19b-4bd2-a72c-7dda03632282",
   "metadata": {},
   "source": [
    "> You can expect to get some warning here, this is ok"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7535b8b-d220-44e8-a56c-97e250c36596",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "ds = load_dataset('MohamedRashad/characters_backstories')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13caeb7f-8a07-4ca2-a770-5b627238c2ac",
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "# Let's take a look at one example\n",
    "ds[\"train\"][400]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dae9106-8015-43da-a6d9-1124dee4bdde",
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "# As this dataset has no validation split, we will create one\n",
    "ds = ds[\"train\"].train_test_split(test_size=0.2, seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea1602d-504b-43de-87ad-fcb35b9e61f7",
   "metadata": {
    "height": 284
   },
   "outputs": [],
   "source": [
    "# We'll create a tokenizer from model checkpoint\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=False)\n",
    "\n",
    "# We'll need padding to have same length sequences in a batch\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Define a tokenization function that first concatenates text and target\n",
    "def tokenize_function(example):\n",
    "    merged = example[\"text\"] + \" \" + example[\"target\"]\n",
    "    batch = tokenizer(merged, padding='max_length', truncation=True, max_length=128)\n",
    "    batch[\"labels\"] = batch[\"input_ids\"].copy()\n",
    "    return batch\n",
    "\n",
    "# Apply it on our dataset, and remove the text columns\n",
    "tokenized_datasets = ds.map(tokenize_function, remove_columns=[\"text\", \"target\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a42417b8-ffa8-4d96-92ea-d8d949d87d5e",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "# Let's check out one prepared example\n",
    "print(tokenizer.decode(tokenized_datasets[\"train\"][900]['input_ids']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8d6b17-a63d-41f1-92cf-416064b52156",
   "metadata": {},
   "source": [
    "### Training\n",
    "Let's finetune a pretrained language model on our dataset using HF Transformers and their wandb integration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f131eb-979e-40f6-9e28-19756beaa8e4",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "# We will train a causal (autoregressive) language model from a pretrained checkpoint\n",
    "model = AutoModelForCausalLM.from_pretrained(model_checkpoint);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7345ab23-8d12-4d4c-a39d-bb2202bff218",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "# Start a new wandb run\n",
    "run = wandb.init(project='dlai_lm_tuning', job_type=\"training\", anonymous=\"allow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d74ee155-3c30-4ef2-9c4d-fd8ee222c50c",
   "metadata": {
    "height": 233
   },
   "outputs": [],
   "source": [
    "# Define training arguments\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "training_args = TrainingArguments(\n",
    "    f\"{model_name}-finetuned-characters-backstories\",\n",
    "    report_to=\"wandb\", # we need one line to track experiments in wandb\n",
    "    num_train_epochs=1,\n",
    "    logging_steps=1,\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=1e-4,\n",
    "    weight_decay=0.01,\n",
    "    no_cuda=True, # force cpu use, will be renamed `use_cpu`\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af62105f-a478-436f-88a2-5c1d78b9d20a",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "# We'll use HF Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"test\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01958a56-c22a-4a27-bc71-41c59fc97f05",
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "# Let's train!\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7911e43f-f4ce-4855-9f68-662438af8d24",
   "metadata": {
    "height": 352
   },
   "outputs": [],
   "source": [
    "transformers.logging.set_verbosity_error() # suppress tokenizer warnings\n",
    "\n",
    "prefix = \"Generate Backstory based on following information Character Name: \"\n",
    "\n",
    "prompts = [\n",
    "    \"Frogger Character Race: Aarakocra Character Class: Ranger Output: \",\n",
    "    \"Smarty Character Race: Aasimar Character Class: Cleric Output: \",\n",
    "    \"Volcano Character Race: Android Character Class: Paladin Output: \",\n",
    "]\n",
    "\n",
    "table = wandb.Table(columns=[\"prompt\", \"generation\"])\n",
    "\n",
    "for prompt in prompts:\n",
    "    input_ids = tokenizer.encode(prefix + prompt, return_tensors=\"pt\")\n",
    "    output = model.generate(input_ids, do_sample=True, max_new_tokens=50, top_p=0.3)\n",
    "    output_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "    table.add_data(prefix + prompt, output_text)\n",
    "    \n",
    "wandb.log({'tiny_generations': table})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7caafaf-dfb7-413d-a329-8520aea5f73a",
   "metadata": {},
   "source": [
    "**Note**: LLM's don't always generate the same results. Your generated characters and backstories may differ from the video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3083c6a3-fdb8-44ab-a028-c0a222a2fdef",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e429a9-5255-4a76-a27b-945c2afd17e4",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfa3bfcc-2885-4eb8-8a18-a236c69e1a98",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
